#include "WVTree.h"
#include <algorithm>

WVNode::WVNode(WVTree* tree, WVNode* parent, unsigned level)
	: tree_(tree)
	, parent_(parent)
	, level_(level)
	, negtive_(nullptr)
	, positive_(nullptr)
{
}

WVNode::~WVNode()
{
	delete negtive_;
	delete positive_;
}

void WVNode::Insert(const WVector& vector)
{
	if (level_ == tree_->GetMaxLevels() - 1)
	{
		vectors_.push_back(vector);
		return;
	}

	const float margin = tree_->GetMargin();
	float axis = vector[level_];

	if (axis <= margin)
	{
		if (!negtive_)
		{
			negtive_ = new WVNode(tree_, this, level_ + 1);
		}
		negtive_->Insert(vector);
	}

	if (axis > -margin)
	{
		if (!positive_)
		{
			positive_ = new WVNode(tree_, this, level_ + 1);
		}
		positive_->Insert(vector);
	}
}

bool WVNode::Find(const WVector& point, float threshold, WVector& result) const
{
	if (level_ == tree_->GetMaxLevels() - 1)
	{
		return FindImpl(point, threshold, result);
	}

	if (point[level_] <= 0.0f)
	{
		if (negtive_)
		{
			return negtive_->Find(point, threshold, result);
		}
	}
	else
	{
		if (positive_)
		{
			return positive_->Find(point, threshold, result);
		}
	}

	return false;
}

bool WVNode::FindAll(const WVector& point, float threshold, std::vector<WVector>& results) const
{
	if(level_ == tree_->GetMaxLevels() - 1)
	{
		return FindAllImpl(point, threshold, results);
	}

	if (point[level_] <= 0.0f)
	{
		if (negtive_)
		{
			return negtive_->FindAll(point, threshold, results);
		}
	}
	else
	{
		if (positive_)
		{
			return positive_->FindAll(point, threshold, results);
		}
	}

	return false;
}

size_t WVNode::GetNumNodes() const
{
	size_t numNodes = 1;

	if (negtive_)
	{
		numNodes += negtive_->GetNumNodes();
	}
	if (positive_)
	{
		numNodes += positive_->GetNumNodes();
	}

	return numNodes;
}

size_t WVNode::GetNumWords() const
{
	size_t numWords = vectors_.size();

	if (negtive_)
	{
		numWords += negtive_->GetNumWords();
	}
	if (positive_)
	{
		numWords += positive_->GetNumWords();
	}

	return numWords;
}

size_t WVNode::GetMaxWords() const
{
	size_t maxPoints = vectors_.size();

	if (negtive_)
	{
		maxPoints = std::max(maxPoints, negtive_->GetMaxWords());
	}
	if (positive_)
	{
		maxPoints = std::max(maxPoints, positive_->GetMaxWords());
	}

	return maxPoints;
}

bool WVNode::FindImpl(const WVector& point, float threshold, WVector& result) const
{
	if (vectors_.empty())
	{
		return false;
	}

	float similar = vectors_[0].CalcSimilar(point);
	size_t index = 0;
	for (size_t i = 1; i < vectors_.size(); ++i)
	{
		float s = vectors_[i].CalcSimilar(point);
		if (similar < s)
		{
			similar = s;
			index = i;
		}
	}

	if (similar > threshold)
	{
		result = vectors_[index];
		return true;
	}

	return false;
}

bool WVNode::FindAllImpl(const WVector& point, float threshold, std::vector<WVector>& results) const
{
	if (vectors_.empty())
	{
		return false;
	}

	for (size_t i = 0; i < vectors_.size(); ++i)
	{
		float similar = vectors_[i].CalcSimilar(point);
		if (similar > threshold)
		{
			results.push_back(vectors_[i]);
		}
	}

	return !results.empty();
}

////////////////////////////////////////////////////////////////////////////////////////////////////

WVTree::WVTree(unsigned maxLevels, float margin) : WVNode(nullptr, nullptr, 0)
, maxLevels_(maxLevels)
, margin_(margin)
{
	tree_ = this;
}

void WVTree::SetVectors(const std::vector<WVector>& vectors)
{
	for (size_t i = 0; i < vectors.size(); ++i)
	{
		Insert(vectors[i]);
	}
}

void WVTree::Report()
{
	printf("Num nodes = %d\n", GetNumNodes());
	printf("Num words = %d\n", GetNumWords());
	printf("Max words = %d, Average words = %d\n", GetMaxWords(), GetNumWords() / GetNumNodes());
}
